package com.realshare.frame.util.safety.china;

import com.realshare.frame.util.safety.Base64Utils;
import com.realshare.frame.util.safety.FileHandler;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.*;
import java.util.Date;

public class Sm4Util extends GmBaseUtil {
    public static final String ALGORITHM_NAME = "SM4";
    public static final String ALGORITHM_NAME_ECB_PADDING = "SM4/ECB/PKCS5Padding";
    public static final String ALGORITHM_NAME_CBC_PADDING = "SM4/CBC/PKCS5Padding";
    public static final int DEFAULT_KEY_SIZE = 128;

    public static String generateKey() {
        return generateKey(DEFAULT_KEY_SIZE);
    }

    public static String generateKey(int keySize) {
        try {
            KeyGenerator kg = KeyGenerator.getInstance(ALGORITHM_NAME, BouncyCastleProvider.PROVIDER_NAME);
            kg.init(keySize, new SecureRandom());
            byte[] bytes = kg.generateKey().getEncoded();
            return Base64Utils.encoder(bytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String encryptEcb(String data, String key) {
        try {
            Cipher cipher = generateEcbCipher(ALGORITHM_NAME_ECB_PADDING, Cipher.ENCRYPT_MODE, key);
            if (cipher != null) {
                byte[] bytes = cipher.doFinal(data.getBytes());
                return Base64Utils.encoder(bytes);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String decryptEcb(String data, String key) {
        try {
            Cipher cipher = generateEcbCipher(ALGORITHM_NAME_ECB_PADDING, Cipher.DECRYPT_MODE, key);
            if (cipher != null) {
                byte[] bytes = Base64Utils.decoder(data);
                return new String(cipher.doFinal(bytes));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return "";
    }

    /**
     * 加密文件
     *
     * @param sourcePath  源文件
     * @param encryptPath 加密文件
     * @param key         密钥
     */
    public static void encryptEcbFile(String sourcePath, String encryptPath, String key) {
        try {
            Cipher enCipher = generateEcbCipher(ALGORITHM_NAME_ECB_PADDING, Cipher.ENCRYPT_MODE, key);
            FileHandler.handlerFile(encryptPath, sourcePath, enCipher);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 解密文件
     *
     * @param encryptPath 加密文件
     * @param decryptPath 解密文件
     * @param key         密钥
     */
    public static void decryptEcbFile(String encryptPath, String decryptPath, String key) {
        try {
            Cipher deCipher = generateEcbCipher(ALGORITHM_NAME_ECB_PADDING, Cipher.DECRYPT_MODE, key);
            FileHandler.handlerFile(decryptPath, encryptPath, deCipher);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static String encryptCbc(String data, String key, String iv) {
        try {
            Cipher cipher = generateCbcCipher(ALGORITHM_NAME_CBC_PADDING, Cipher.ENCRYPT_MODE, key, iv);
            if (cipher != null) {
                byte[] bytes = cipher.doFinal(data.getBytes());
                return Base64Utils.encoder(bytes);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return "";
    }

    public static String decryptCbc(String data, String key, String iv) {
        try {
            Cipher cipher = generateCbcCipher(ALGORITHM_NAME_CBC_PADDING, Cipher.DECRYPT_MODE, key, iv);
            if (cipher != null) {
                byte[] bytes = Base64Utils.decoder(data);
                return new String(cipher.doFinal(bytes));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return "";
    }

    /**
     * 加密文件
     *
     * @param sourcePath  源文件
     * @param encryptPath 加密文件
     * @param key         密钥
     */
    public static void encryptCbcFile(String sourcePath, String encryptPath, String key, String iv) {
        try {
            Cipher enCipher = generateCbcCipher(ALGORITHM_NAME_CBC_PADDING, Cipher.ENCRYPT_MODE, key, iv);
            FileHandler.handlerFile(encryptPath, sourcePath, enCipher);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 解密文件
     *
     * @param encryptPath 加密文件
     * @param decryptPath 解密文件
     * @param key         密钥
     */
    public static void decryptCbcFile(String encryptPath, String decryptPath, String key, String iv) {
        try {
            Cipher deCipher = generateCbcCipher(ALGORITHM_NAME_CBC_PADDING, Cipher.DECRYPT_MODE, key, iv);
            FileHandler.handlerFile(decryptPath, encryptPath, deCipher);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static Cipher generateEcbCipher(String algorithmName, int mode, String key) {
        try {
            byte[] keyBytes = Base64Utils.decoder(key);
            Cipher cipher = Cipher.getInstance(algorithmName, BouncyCastleProvider.PROVIDER_NAME);
            Key sm4Key = new SecretKeySpec(keyBytes, ALGORITHM_NAME);
            cipher.init(mode, sm4Key);
            return cipher;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * @param algorithmName
     * @param mode
     * @param key
     * @param iv            初始化向量，随意填充，必须16长度
     * @return Cipher
     */
    private static Cipher generateCbcCipher(String algorithmName, int mode, String key, String iv) {
        try {
            byte[] tempBytes = iv.getBytes();
            byte[] ivBytes = new byte[16];
            for (int i = 0; i < tempBytes.length && i < ivBytes.length; i++) {
                ivBytes[i] = tempBytes[i];
            }
            byte[] keyBytes = Base64Utils.decoder(key);
            Cipher cipher = Cipher.getInstance(algorithmName, BouncyCastleProvider.PROVIDER_NAME);
            Key sm4Key = new SecretKeySpec(keyBytes, ALGORITHM_NAME);
            IvParameterSpec ivParameterSpec = new IvParameterSpec(ivBytes);
            cipher.init(mode, sm4Key, ivParameterSpec);
            return cipher;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static void main(String[] args) {
        String key = generateKey();
        String source = "一醉汉不慎从三楼掉下，引来路人围观，一警察过来：发生了什么事？醉汉：不清楚，我也是刚到！";
        // 初始化向量，随意填充，必须16长度
        String iv = "1234asd000000000";
        Date date = new Date();
        testEcbFile(source, key);
        testEcb(source, key);
        testCbcFile(source, key, iv);
        testCbc(source, key, iv);
        System.out.println("加密耗时：" + (new Date().getTime() - date.getTime()) + "毫秒");
    }

    private static void testEcbFile(String source, String key) {
        String sourcePath = "D:\\1.txt";
        String encryptPath = "D:\\sm4Ecb_encrypt.txt";
        String decryptPath = "D:\\sm4Ecb_decrypt.txt";
        encryptEcbFile(sourcePath, encryptPath, key);
        decryptEcbFile(encryptPath, decryptPath, key);
    }

    private static void testEcb(String source, String key) {
        System.out.println("原文字：" + source);
        System.out.println("密钥：" + key);
        String encryptStr = encryptEcb(source, key);
        System.out.println("sm4Ecb加密后：" + encryptStr);
        String decryptStr = decryptEcb(encryptStr, key);
        System.out.println("sm4Ecb解密后：" + decryptStr);
    }

    private static void testCbcFile(String source, String key, String iv) {
        String sourcePath = "D:\\1.txt";
        String encryptPath = "D:\\sm4Cbc_encrypt.txt";
        String decryptPath = "D:\\sm4Cbc_decrypt.txt";
        encryptCbcFile(sourcePath, encryptPath, key, iv);
        decryptCbcFile(encryptPath, decryptPath, key, iv);
    }

    private static void testCbc(String source, String key, String iv) {
        System.out.println("原文字：" + source);
        System.out.println("密钥：" + key);
        String encryptStr = encryptCbc(source, key, iv);
        System.out.println("sm4Cbc加密后：" + encryptStr);
        String decryptStr = decryptCbc(encryptStr, key, iv);
        System.out.println("sm4Cbc解密后：" + decryptStr);
    }
}
